
# performance benchmarking

def benchmark(func, **kwargs):
	import time, torch
	from tqdm import tqdm
	var_name, var_vals = None, None
	for k,v in kwargs.items():
		if isinstance(v, (tuple, list)):
			var_name, var_vals = k, v
	flops = []
	print(var_name, var_vals)
	for var_val in tqdm(var_vals):
		kwargs[var_name] = var_val
		ts = time.time()
		flop = func(**kwargs)
		te = time.time()
		t = te-ts
		flops.append(flop/t)

	# flops
	X = torch.Tensor(var_vals)
	Y = torch.Tensor([]).unsqueeze(0)
	Y = torch.Tensor(flops).unsqueeze(0)
	plot_(X, Y, xlabel=var_name, ylabel='flops',
		filename=f'{var_name}-flops.pdf')

	# scaling
	Y = torch.Tensor([(f/flops[0]) 
					for f in flops]).unsqueeze(0)
	plot_(X, Y, xlabel=var_name, ylabel='scaling rate', 
		filename=f'{var_name}-scaling.pdf')

	if var_name=='cores':
		# parallel efficiency
		Y = torch.Tensor([(f/c)/(flops[0]/var_vals[0])
			for f,c in zip(flops,var_vals)]).unsqueeze(0)
		plot_(X, Y, xlabel=var_name, 
			ylabel='parallel efficiency', 
			filename=f'{var_name}-efficiency.pdf')


import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

def plot_(x_data, y_data, y_err=None, legends=None, title=None, xlabel=None, ylabel=None, filename=None):
	with plt.style.context(('fivethirtyeight')):
		fig, ax = plt.subplots(squeeze=True)
		for i in range(y_data.size(0)):
			ax.plot(x_data.numpy(), y_data[i].numpy(), label=legends[i] if legends else '?')
			if (not y_err is None) and (not y_err[i] is None):
				ax.fill_between(x_data.numpy(), (y_data[i]-y_err[i]).numpy(), (y_data[i]+y_err[i]).numpy(), alpha=0.3)
		if legends: ax.legend(loc=2)
		if title: ax.set_title(title)
		if xlabel: ax.set_xlabel(xlabel)
		if ylabel: ax.set_ylabel(ylabel)
		fig.tight_layout()
	if filename:
		plt.savefig(filename)
		plt.close()
	else:
		plt.show()
		return plt